from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI


def get_schema(_):
    return db.get_table_info()

def run_query(query):
    return db.run(query)

model = ChatOpenAI(
    openai_api_key="key",
    openai_api_base="https://api.moonshot.cn/v1",
    model="moonshot-v1-8k",
    temperature=0,
    request_timeout=60,
    max_retries=3,
)

template = """根据下面的表模式，编写一个SQL查询来回答用户的问题：
{schema}

问题：{question}
SQL查询："""
prompt = ChatPromptTemplate.from_template(template)

db = SQLDatabase.from_uri("sqlite:///./Chinook.db")
sql_response = (
    RunnablePassthrough.assign(schema=get_schema)
    | prompt
    | model.bind(stop=["\nSQLResult:"])
    | StrOutputParser()
)
sql_response.invoke({"question": "有多少员工？"})

template = """根据下面的表模式，问题，SQL查询和SQL响应，编写一个自然语言回答：
{schema}

问题：{question}
SQL查询：{query}
SQL响应：{response}"""

prompt_response = ChatPromptTemplate.from_template(template)
full_chain = (
    RunnablePassthrough.assign(query=sql_response).assign(
        schema=get_schema,
        response=lambda x: db.run(x["query"]),
    )
    | prompt_response
    | model
)
full_chain.invoke({"question": "有多少员工？"})
